# -*- coding: utf-8 -*-
"""
Created on Fri Jul  2 18:57:54 2021

@author: MA
"""

import numpy as np
import os
import nibabel as nib
join = os.path.join
from skimage import measure,morphology


#%% refine hard cases
hard_seg_path = 'path to/FindVertebrae/codes/seg_temp/T362-SegHard'
ensemble_path = 'path to/FindVertebrae/codes/seg_temp/seg-Ensemble'
save_path = 'path to/FindVertebrae/segmentation_results'
names = sorted(os.listdir(ensemble_path))
# fill in small holes
def fill_small_holes(ensemble_data):
    ensemble_data = np.uint8(ensemble_data)
    refine_data = np.zeros_like(ensemble_data, dtype=np.uint8)  
    classes = [2,3,4,5,6,7,10,11,12,13,14,15,16,17,18,19] # do not refine 1, 8, 9 because they are in the top and down
    for i in range(ensemble_data.shape[-1]):
        for c in classes:# ignore background
            if np.sum(ensemble_data[:,:,i]==c)>0:
                # print(i, c)
                slice_i_c = morphology.remove_small_holes(ensemble_data[:,:,i]==c, 200)
                refine_data[:,:,i][slice_i_c>0] = c
    refine_data[ensemble_data==1] = 1
    refine_data[ensemble_data==8] = 8
    refine_data[ensemble_data==9] = 9
    return refine_data.astype(np.uint8)

#%% add label 18 based on hard model      
name_processed = []
for name in names:
    hard_nii = nib.load(join(hard_seg_path, name))
    hard_data = hard_nii.get_fdata()
    if len(np.unique(hard_data)) == 2:
        ensemble_nii = nib.load(join(ensemble_path, name))
        ensemble_data = ensemble_nii.get_fdata()
        if np.sum(ensemble_data[hard_data==3])<5: # label 18; T10/T11
            print(name, 'hard cases lack of label 18', np.unique(hard_data))
            ensemble_data[hard_data==3] = 18
            ensemble_data = fill_small_holes(ensemble_data)
            save_nii = nib.Nifti1Image(ensemble_data.astype(np.uint8), ensemble_nii.affine, ensemble_nii.header)
            nib.save(save_nii, join(save_path, name))
            name_processed.append(name)
            # names.remove(name)
for name in name_processed:
    names.remove(name)     
#%% check the middle slice [440,:,:]; identify obvious wrong slices
sanity_slice = np.uint32(ensemble_data.shape[0]/2) # middle slice
sanity_labels = [5,6,15,16]
for name in names:
    ensemble_nii = nib.load(join(ensemble_path, name))
    ensemble_data = ensemble_nii.get_fdata()
    region_num = 0
    for label_id in sanity_labels:
        region_num+= np.max(measure.label(ensemble_data[sanity_slice,:,:]==label_id))
    if region_num>8:
        print(name, 'obvious wrong case!')
        seg1 = nib.load(join('./seg_temp/seg1', name)).get_fdata()
        seg2 = nib.load(join('./seg_temp/seg2', name)).get_fdata()
        seg3 = nib.load(join('./.seg_temp/seg3', name)).get_fdata()
        seg4 = nib.load(join('./seg_temp/seg4', name)).get_fdata()
        seg_3d = nib.load(join('./seg_temp/seg-3D', name)).get_fdata()
        seg_stack = np.uint8(np.stack((ensemble_data, seg1, seg2, seg3, seg4),axis=0))
        seg_common = np.prod(seg_stack, axis=0) # ensemble_data*seg1*seg2*seg3*seg4
        classes = np.unique(ensemble_data)
        slices_num = ensemble_data.shape[-1]
        case_reg_num = np.zeros((slices_num, seg_stack.shape[0]))
        seg_stack_roi = np.zeros_like(seg_stack, dtype=np.uint8)
        for seg_id in range(seg_stack.shape[0]):
            seg_stack_roi[seg_id,...] = seg_stack[seg_id,...] * np.uint8(seg_common>0) # extract common seg results
            for i in range(slices_num):
                for j in classes:
                    case_reg_num[i, seg_id] += np.max(measure.label(seg_stack_roi[seg_id,:,:,i]==j))

        good_ids = np.argmin(case_reg_num,axis=-1)
        final_seg = np.zeros_like(ensemble_data,dtype=np.uint8)
        for i in range(slices_num):
                final_seg[:,:,i]= seg_stack[good_ids[i],:,:,i]
        final_seg[:,:,np.uint8(ensemble_data.shape[-1]/2-1)] = seg_3d[:,:,np.uint8(ensemble_data.shape[-1]/2-1)] # replace middle slice with 3D seg
        finel_seg = fill_small_holes(final_seg)
        save_nii = nib.Nifti1Image(final_seg.astype(np.uint8), ensemble_nii.affine, ensemble_nii.header)
        nib.save(save_nii, join(save_path, name))
    else:
        final_seg = fill_small_holes(ensemble_data)
        save_nii = nib.Nifti1Image(final_seg.astype(np.uint8), ensemble_nii.affine, ensemble_nii.header)
        nib.save(save_nii, join(save_path, name))

        


            
        








